import sys
import puzautils.table.puza_memory_table as pmt
import puzautils.table.puza_row as pr
import puzautils.table.puza_stream_table as pst
from datetime import datetime
from time import mktime


input_filename = sys.argv[1]
output_filename = ''
if len(sys.argv) > 2:
    output_filename = sys.argv[2]
    

class LocalFriendship:
    def filter(self, s):
        try:
            r = pr.PuzaRow.Load(s)
            print r['id'], r['followers_count'], r['friends_count'], r['statuses_count']
        except:
            pass
    def post(self):
        pass

class TweetIDUserID:
    def __init__(self):
        global output_filename
        self.OutputTable = pst.PuzaStreamTable()
        self.OutputTable.Save(output_filename)
        
    def filter(self, s):
        try:
            r = pr.PuzaRow.Load(s)
            new_r = pr.PuzaRow()
            new_r['id'] = r['id']
            new_r['user_id'] = r['user']['id']
            self.OutputTable.StreamOutRow(new_r, pointer = 0, lformat = 1, block_number = 1024)
        except:
            pass

    def post(self):
        self.OutputTable.StreamClose()

class TweetUserPreTweetUser:
    def __init__(self):
        global output_filename
        self.OutputTable = pst.PuzaStreamTable()
        self.OutputTable.Save(output_filename)
        
    def filter(self, s):
        try:
            r = pr.PuzaRow.Load(s)
            new_row = pr.PuzaRow()
            new_row['id'] = r['id']
            new_row['user_id'] = r['user']['id']
            if r['retweeted_status'] != None:
                new_row['previous_id'] = r['retweeted_status']['id']
                new_row['previous_user_id'] = r['retweeted_status']['user']['id']
            self.OutputTable.StreamOutRow(new_row, pointer = 0, lformat = 1, block_number = 1024)
        except:
            pass

    def post(self):
        self.OutputTable.StreamClose()

class DistinctField:
    
    def __init__(self):
        self._pool = set()

    def set_key(self,field_name):
        self.field_name = field_name
        
    def filter(self, s):
        try:
            r = pr.PuzaRow.Load(s)
            item = r[self.field_name]
            self._pool.add(item)
        except:
            pass

    def post(self):
        print self.field_name, len(self._pool)
    

## ===============
if __name__ == '__main__':
    A = pst.PuzaStreamTable()
    A.Read(input_filename)
    # L = LocalFriendship()
    # L = TweetIDUserID()
    # L = TweetUserPreTweetUser()
    L = DistinctField()
    L.set_key('user_id')
    A.ReadLines(L.filter)
    L.post()
